package com.example;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.StandardSocketOptions;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * @author linfeng
 * @date 2021/12/27 9:30
 */
public class NonBlockSocketServer {

    private static Map<SocketChannel, ByteBuffer> writeBufferMap = new HashMap<>();

    private static boolean haveDataToBeWrite(SocketChannel socketChannel) {
        if (writeBufferMap.containsKey(socketChannel)) {
            ByteBuffer byteBuffer = writeBufferMap.getOrDefault(socketChannel, null);
            if (byteBuffer != null && byteBuffer.hasRemaining()) {
                return true;
            }
        }
        return false;
    }

    private static void processChannelRead(SocketChannel socketChannel) throws IOException {
        ByteBuffer buffer = ByteBuffer.allocate(1024);
        while (true) {
            int readBytes = socketChannel.read(buffer);
            if (readBytes <= 0) break;
            buffer.flip();
            while (buffer.hasRemaining()) {
                System.out.print((char) buffer.get());
            }
        }
        ByteBuffer buffer2 = ByteBuffer.allocate("hello, world.".length() * 1000000);
        writeBufferMap.put(socketChannel, buffer2);
        buffer2.put("hello, world.".repeat(1000000).getBytes(StandardCharsets.UTF_8));
        buffer2.flip();
        socketChannel.write(buffer2);
        if (buffer2.hasRemaining()) {
            writeBufferMap.put(socketChannel, buffer2);
        }
    }

    private static void processChannelWrite(SocketChannel socketChannel) throws IOException {
        ByteBuffer buffer2 = writeBufferMap.get(socketChannel);
        if (buffer2 == null || !socketChannel.isOpen()) {
            return;
        }
        socketChannel.write(buffer2);
        if (!buffer2.hasRemaining()) {
            writeBufferMap.remove(socketChannel);
        }
    }

    public static void main(String[] args) throws IOException {
        Selector selector = Selector.open();
        ServerSocketChannel serverSocketChannel = ServerSocketChannel.open();
        serverSocketChannel.configureBlocking(false);
        serverSocketChannel.bind(new InetSocketAddress("0.0.0.0", 8088));
        serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT);
        while (true) {
            int readyChannels = selector.select();
            if (readyChannels == 0) {
                continue;
            }
            Iterator<SelectionKey> selectionKeyIterator = selector.selectedKeys().iterator();
            while (selectionKeyIterator.hasNext()) {
                SelectionKey selectedKey = selectionKeyIterator.next();
                if (selectedKey.isAcceptable()) {
                    ServerSocketChannel channel = (ServerSocketChannel) selectedKey.channel();
                    SocketChannel socketChannel = channel.accept();
                    socketChannel.configureBlocking(false);
                    // 设置发送缓冲区大小
                    socketChannel.setOption(StandardSocketOptions.SO_SNDBUF, 128);
                    socketChannel.register(selector, SelectionKey.OP_READ | SelectionKey.OP_WRITE);
                    selectionKeyIterator.remove();
                    continue;
                }
                SocketChannel channel = (SocketChannel) selectedKey.channel();
                if (selectedKey.isReadable()) {
                    processChannelRead(channel);
                }
                if (selectedKey.isWritable() && haveDataToBeWrite(channel)) {
                    processChannelWrite(channel);
                }
                selectionKeyIterator.remove();
            }
        }
    }
}
